import os
import glob
import xml.etree.ElementTree as ET

def convert_voc_to_yolo(xml_file, class_names):
    """
    将VOC格式的XML标注转换为YOLO格式的txt文件
    
    Args:
        xml_file: XML文件路径
        class_names: 类别名称列表
    """
    tree = ET.parse(xml_file)
    root = tree.getroot()
    
    # 获取图像尺寸
    size = root.find('size')
    width = float(size.find('width').text)
    height = float(size.find('height').text)
    
    # 创建与XML文件同名的txt文件
    txt_file = os.path.splitext(xml_file)[0] + '.txt'
    
    with open(txt_file, 'w', encoding='utf-8') as f:
        # 处理每个目标
        for obj in root.findall('object'):
            # 获取类别名称
            class_name = obj.find('name').text
            
            # 如果类别不在列表中，跳过
            if class_name not in class_names:
                print(f"警告: 未知类别 '{class_name}' 在文件 {xml_file}")
                continue
            
            # 获取类别索引
            class_idx = class_names.index(class_name)
            
            # 获取边界框坐标
            bbox = obj.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)
            
            # 转换为YOLO格式（中心点坐标和宽高，归一化到0-1）
            x_center = ((xmin + xmax) / 2) / width
            y_center = ((ymin + ymax) / 2) / height
            bbox_width = (xmax - xmin) / width
            bbox_height = (ymax - ymin) / height
            
            # 写入txt文件
            f.write(f"{class_idx} {x_center:.6f} {y_center:.6f} {bbox_width:.6f} {bbox_height:.6f}\n")

def main():
    # 当前目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 创建images和labels目录
    images_dir = os.path.join(current_dir, 'images')
    labels_dir = os.path.join(current_dir, 'labels')
    
    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(labels_dir, exist_ok=True)
    
    # 从dataset.yaml中读取类别名称
    yaml_file = os.path.join(current_dir, 'dataset.yaml')
    class_names = []
    
    if os.path.exists(yaml_file):
        with open(yaml_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            started = False
            for line in lines:
                if line.strip().startswith('names:'):
                    started = True
                    continue
                if started and line.strip():
                    if ':' in line:
                        parts = line.strip().split(':')
                        if len(parts) >= 2:
                            class_name = parts[1].strip()
                            class_names.append(class_name)
    
    if not class_names:
        print("错误: 无法从dataset.yaml中读取类别名称")
        return
    
    print(f"找到{len(class_names)}个类别: {class_names}")
    
    # 获取所有XML文件
    xml_files = glob.glob(os.path.join(current_dir, '*.xml'))
    print(f"找到{len(xml_files)}个XML文件")
    
    # 获取所有图像文件
    image_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp']
    image_files = []
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(current_dir, f'*{ext}')))
    
    print(f"找到{len(image_files)}个图像文件")
    
    # 转换所有XML文件
    for xml_file in xml_files:
        try:
            convert_voc_to_yolo(xml_file, class_names)
            print(f"转换完成: {xml_file}")
        except Exception as e:
            print(f"转换失败: {xml_file}, 错误: {str(e)}")
    
    # 移动所有图像文件到images目录
    for img_file in image_files:
        basename = os.path.basename(img_file)
        dest_path = os.path.join(images_dir, basename)
        try:
            # 复制图像文件，而不是移动，以防出错
            import shutil
            shutil.copy2(img_file, dest_path)
            print(f"复制图像: {basename}")
        except Exception as e:
            print(f"复制图像失败: {img_file}, 错误: {str(e)}")
    
    # 移动所有txt文件到labels目录
    txt_files = glob.glob(os.path.join(current_dir, '*.txt'))
    for txt_file in txt_files:
        # 跳过dataset.yaml文件
        if os.path.basename(txt_file) == 'dataset.yaml':
            continue
        
        basename = os.path.basename(txt_file)
        dest_path = os.path.join(labels_dir, basename)
        try:
            # 复制txt文件，而不是移动，以防出错
            import shutil
            shutil.copy2(txt_file, dest_path)
            print(f"复制标签: {basename}")
        except Exception as e:
            print(f"复制标签失败: {txt_file}, 错误: {str(e)}")
    
    print("\n转换完成!")
    print(f"图像文件已复制到: {images_dir}")
    print(f"标签文件已复制到: {labels_dir}")
    print("现在可以运行评估脚本了。")

if __name__ == "__main__":
    main() 